package array.prefixsum;

import java.util.HashMap;

public class SubArrayDivByK {

    public int subarrayDivByK(int[] nums, int k) {

        int n = nums.length;
        // nums 的前缀和数组，注意索引偏移，preSum[i]的值为sum(nums[0]...nums[i-1])
        int[] preSum = new int[n + 1];
        preSum[0] = 0;
        // 前缀和余数的计算，索引偏移，方便快速查找所需的前缀和余数的数量
        HashMap<Integer, Integer> remainderToCount = new HashMap<>();
        preSum[0] = 0;
        remainderToCount.put(0, 1);
        int res = 0;
        // 计算 nums 的前缀和余数并更新计数器
        for (int i = 1; i <= n; i++) {
            preSum[i] = preSum[i - 1] + nums[i - 1];
            int remainder = preSum[i] % k;
            if (remainder < 0) {
                // 考虑到 preSum[i] 可能是负数，根据Java 求模的特性，-2 % 5 = -2
                // 但是实际想要正余数3，所以这里需要对 remainder 进行修正
                remainder += k;
            }
            // 看看之前 nums[0...i-1] 中是否也存在前缀和余数为 remainder 的子数组
            if (remainderToCount.containsKey(remainder)) {
                // 如果存在，则说明找到了可以整除 k 的子数组，累加子数组数量
                res += remainderToCount.get(remainder);
                remainderToCount.put(remainder, remainderToCount.get(remainder) + 1);
            } else {
                // 如果不存在，那么nums[0...i] 是第一个前缀和余数为 remainder 的子数组
                remainderToCount.put(remainder, 1);
            }
        }

        return res;
    }
}
